package com.feidee.fd.sml.algorithm.component.ml.regression

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.regression.GeneralizedLinearRegression

/**
  * @Author songhaicheng
  * @Date 2019/3/20 18:36
  * @Description
  * @Reviewer dongguosheng
  */
case class GLRegressionParam (
                               override val input_pt: String,
                               override val output_pt: String,
                               override val hive_table: String,
                               override val flow_time: String,
                               override val featuresCol: String,
                               override var labelCol: String,
                               override var predictionCol: String,
                               override val modelPath: String,
                               override val metrics: Array[String],
                               // 样本权重列
                               weightCol: String,
                               // 连接函数预测值，不填则不输出
                               linkPredictionCol: String,
                               // 连接函数，支持 [identity, log, inverse, logit, probit, cloglog, sqrt]
                               link: String,
                               // 模型中使用的误差分布类型，支持 [gaussian, binomial, poisson, gamma, tweedie], 默认 gaussian（不同分布支持的连接函数有所不同）
                               family: String,
                               // 优化方法，目前仅支持 irls
                               solver: String,
                               // 最大迭代次数，>= 0，默认 25
                               maxIter: Int,
                               // 连接指数，仅在误差分布类型为 "tweedie" 时有用，默认值为 1 - 误差指数（0, 1, -1, 0.5 时连接函数分别为 log, identity, inverse, sqrt）
                               linkPower: Double,
                               // 正则化系数，>= 0，默认 0.0
                               regParam: Double,
                               // 收敛系数，>= 0，默认 1E-6
                               tol: Double,
                               // 误差指数，仅在误差分布类型为 默认 0.0 "tweedie" 时有用，= 0 或 >= 1，默认 0（0, 1, 2 的分布类型分别对应 gaussian、poisson、gamma）
                               variancePower: Double,
                               // 是否需要计算截距，默认 true
                               fitIntercept: Boolean
                             ) extends MLParam {
  def this() = this(null, null, null, null, "features", "label", "prediction", null, new Array[String](0),
    null, null, null, "gaussian", "irls", 25, Double.NaN, 0.0, 1E-6, 0.0, true)

  override def verify(): Unit = {
    super.verify()
    val links = Array("identity", "log", "inverse", "logit", "probit", "cloglog", "sqrt")
    val families = Array("gaussian", "binomial", "poisson", "gamma", "tweedie")
    val solvers = Array("irls")

    require(!tool.isNotNull(link) ||
      links.contains(link.toLowerCase),
      s"param link only accepts ${links.mkString("[", ", ", "]")}, but has $link")
    require(families.contains(family.toLowerCase),
      s"param family only accepts ${families.mkString("[", ", ", "]")}, but has $family")
    require(solvers.contains(solver.toLowerCase), s"param solver only accepts ${solvers.mkString("[", ", ", "]")}," +
      s" but has $solver")
    require(maxIter >= 0, "param maxIter can't be negative")
    require(regParam >= 0, "param regParam can't be negative")
    require(tol >= 0, "param tol can't be negative")
    require(variancePower == 0 || variancePower >= 1, "param variancePower must be equal to 0 or greater than 1")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("weightCol" -> weightCol)
    map += ("linkPredictionCol" -> linkPredictionCol)
    map += ("link" -> link)
    map += ("family" -> family)
    map += ("solver" -> solver)
    map += ("maxIter" -> maxIter)
    map += ("linkPower" -> linkPower)
    map += ("regParam" -> regParam)
    map += ("tol" -> tol)
    map += ("variancePower" -> variancePower)
    map += ("fitIntercept" -> fitIntercept)
    map
  }
}


class GLRegressionComponent extends AbstractRegressionComponent[GLRegressionParam] {

  override def setUp(param: GLRegressionParam): PipelineStage = {
    val glr = new GeneralizedLinearRegression()
      .setLabelCol(param.labelCol)
      .setFeaturesCol(param.featuresCol)
      .setPredictionCol(param.predictionCol)
      .setFamily(param.family)
      .setSolver(param.solver)
      .setMaxIter(param.maxIter)
      .setRegParam(param.regParam)
      .setTol(param.tol)
      .setVariancePower(param.variancePower)
      .setFitIntercept(param.fitIntercept)


    if (tool.isNotNull(param.weightCol)) {
      glr.setWeightCol(param.weightCol)
    }
    if (tool.isNotNull(param.linkPredictionCol)) {
      glr.setLinkPredictionCol(param.linkPredictionCol)
    }
    if (tool.isNotNull(param.link)) {
      glr.setLink(param.link)
    }
    if (!param.linkPower.isNaN) {
      glr.setLinkPower(param.linkPower)
    }

    glr
  }

}

object GLRegressionComponent {
  def apply(paramStr: String): Unit = {
    new GLRegressionComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    GLRegressionComponent(args(0))
  }
}
